from collections import deque, namedtuple
import numpy as np
import torch


class PrioritizedReplayBuffer:
    def __init__(self, hyper, device):
        capacity = 1
        while capacity < hyper['buffer_size']:
            capacity *= 2
        self.sumtree = SumTree(capacity)
        self.alpha = hyper['alpha_per']
        self.beta = hyper['beta_per']  # importance-sampling, from initial value increasing to 1
        self.beta_increment = 0.001
        self.epsilon = 0.01
        self.max_td_error = 1  # clipped abs error
        self.experience = namedtuple('Experience', ['state', 'action', 'reward', 'next_state', 'mask'])
        self.batch_size = hyper['batch_size']
        self.device = torch.device(device)

    def __len__(self):
        return self.sumtree.data_num

    def store(self, state, action, reward, next_state, mask):
        transition = self.experience(state, action, reward, next_state, mask)
        max_p = np.max(self.sumtree.tree[-self.sumtree.capacity:])
        if max_p == 0:
            max_p = self.max_td_error
        self.sumtree.add(max_p, transition)  # set the max p for new p

    def sample(self):
        n = self.batch_size
        is_idx, experiences, is_weight = np.empty(n, dtype=np.int32), np.empty(n, dtype=object), np.empty(n)
        pri_seg = self.sumtree.total_p() / n
        self.beta = np.min([1, self.beta + self.beta_increment])
        beg = -self.sumtree.capacity
        end = (self.sumtree.data_num - self.sumtree.capacity) if self.sumtree.data_num < self.sumtree.capacity else -1
        min_prob = np.min(self.sumtree.tree[beg:end]) / self.sumtree.total_p()  # for later calculate ISweight
        for i in range(n):
            a, b = pri_seg * i, pri_seg * (i + 1)
            v = np.random.uniform(a, b)
            idx, p, data = self.sumtree.get_leaf(v)
            prob = p / self.sumtree.total_p()
            is_weight[i] = np.power(prob / min_prob, -self.beta)
            # is_weight = (N*Pj)^(-beta) / maxi_wi
            #          = (N*Pj)^(-beta) / maxi[ (N*Pi)^(-beta) ]
            #          = (N*Pj)^(-beta) / mini[ (N*Pi) ]^(-beta)
            #          = (Pj / mini[Pi])^(-beta)
            is_idx[i] = idx
            experiences[i] = data
        states = torch.tensor(np.vstack([e.state for e in experiences])).float().to(self.device)
        actions = torch.tensor(np.vstack([e.action for e in experiences])).long().to(self.device)
        rewards = torch.tensor(np.vstack([e.reward for e in experiences])).float().to(self.device)
        next_states = torch.tensor(np.vstack([e.next_state for e in experiences])).float().to(self.device)
        masks = torch.tensor(np.vstack([e.mask for e in experiences])).float().to(self.device)
        is_weight = torch.tensor(is_weight).float().to(self.device)
        return states, actions, rewards, next_states, masks, is_idx, is_weight

    def update_td_errors_for_per(self, tree_idx, td_errors):
        abs_errors = np.abs(td_errors) + self.epsilon  # convert to abs and avoid 0
        clipped_errors = np.minimum(abs_errors, self.max_td_error)
        ps = clipped_errors ** self.alpha
        for ti, p in zip(tree_idx, ps):
            self.sumtree.update(ti, p)


class SumTree:
    def __init__(self, capacity):
        self.capacity = capacity
        self.data_pointer = 0
        self.data_num = 0
        self.tree = np.zeros(2 * capacity - 1)  # for all priority values
        # [--------------Parent nodes-------------][-------leaves to recode priority-------]
        #             size: capacity - 1                       size: capacity
        self.data = np.zeros(capacity, dtype=object)  # for all transitions
        # [--------------data frame---------------]
        #             size: capacity

    def add(self, p, data):
        tree_idx = self.data_pointer + self.capacity - 1
        self.data[self.data_pointer] = data  # update data_frame
        self.update(tree_idx, p)
        self.data_pointer += 1
        self.data_num += 1
        if self.data_pointer >= self.capacity:
            self.data_pointer = 0
        if self.data_num >= self.capacity:
            self.data_num = self.capacity

    def update(self, tree_idx, p):
        change = p - self.tree[tree_idx]
        self.tree[tree_idx] = p
        # then propagate the change through tree
        while tree_idx != 0:
            tree_idx = (tree_idx - 1) // 2
            self.tree[tree_idx] += change

    def get_leaf(self, v):
        parent_idx = 0
        while True:
            cl_idx = 2 * parent_idx + 1  # this leaf's left and right kids
            cr_idx = cl_idx + 1
            if cl_idx >= len(self.tree):  # reach bottom, end search
                leaf_idx = parent_idx
                break
            else:  # downward search, always search for a higher priority node
                if v <= self.tree[cl_idx]:
                    parent_idx = cl_idx
                else:
                    v -= self.tree[cl_idx]
                    parent_idx = cr_idx
        data_idx = leaf_idx - (self.capacity - 1)
        return leaf_idx, self.tree[leaf_idx], self.data[data_idx]

    def total_p(self):
        return self.tree[0]  # the root


